import networkx as nx
import matplotlib.pyplot as plt
from itertools import islice


def build_hybrid_topology_shared_tor(k):  # 参数化flat-tree的pods数目
    """构建共享ToR的混合拓扑"""
    # 电网络设置为胖数结构,光网络参考topoopt
    G = nx.Graph()

    # 设备前缀
    SERVER_PREFIX = "Server" # 代替接入交换机的命名
    TOR_PREFIX = "ToR" #
    EDGE_ELECTRIC_PREFIX = "Edge_E"  # 定义一个接入交换机的命名规则
    AGGR_ELECTRIC_PREFIX = "Aggr_E"  # 代替聚合交换机的命名规则
    CORE_ELECTRIC_PREFIX = "Core_E" # 代替核心交换机
    OPTICAL_SWITCH_PREFIX = "Optical_SW"
    # AGGR_OPTICAL_PREFIX = "Aggr_Optical"

    # ================== 共享ToR架构 ==================
    servers_per_tor = 1
    gpus_per_server = 8
    optical_ports_per_tor = k*k*k//4  # 每个ToR的光端口数
    electric_ports_per_tor = 1  # 每个ToR的电端口数

    # 创建ToR交换机和服务器
    for tor_id in range(1, k*k*k//4 + 1):
        # 创建ToR节点（同时支持电和光）
        tor_node = f"ToR{tor_id}"  # 每个ToR交换机命名为ToR1
        G.add_node(tor_node, type='tor',
                   electric_capacity=40, # 每个ToR的电容量
                   optical_capacity=800, #每个ToR的光端口容量
                   electric_ports=electric_ports_per_tor, # 电端口数量
                   optical_ports=optical_ports_per_tor) # 光端口数量

        # 连接服务器
        for s in range(1, servers_per_tor + 1):
            server_node = f"{SERVER_PREFIX}{tor_id}-{s}" #定义服务器的命名规则
            G.add_node(server_node, type='server', gpus=gpus_per_server)
            G.add_edge(tor_node, server_node,
                       type='electrical', # 标记为电端口
                       capacity=10,  # 服务器到ToR的带宽
                       port_type='server')

    # ================== 电网络部分 ==================
    pods=k
    tor_per_pod = k*k // 2  # 每个Pod的ToR数量
    edge_pre_pod=k //2
    total_edges=pods*edge_pre_pod # 总的边缘交换机的数量
    # 创建电网络核心层
    core_e_nodes = [f"{CORE_ELECTRIC_PREFIX}{i}" for i in range(1, int((k//2.0)*(k//2.0)+1))] # 定义核心交换机的命名规则
    G.add_nodes_from(core_e_nodes, type='core_electric', capacity=400) # 核心交换机的容量设为400(电网络）

    # # 创建电网络聚合层
    # aggr_e_nodes = [f"{AGGR_ELECTRIC_PREFIX}{i}" for i in range(1, int((k//2.0*k)+1))] #定义聚合交换机的命名规则
    # G.add_nodes_from(aggr_e_nodes, type='aggregation_electric', capacity=100) # 聚合交换机的容量设为100
    # 
    # # 创建电网络接入层
    # edge_e_nodes=[f"{EDGE_ELECTRIC_PREFIX}{i}" for i in range(1, int((k//2.0*k)+1))]
    # G.add_nodes_from(edge_e_nodes, type='edge_electric', capacity=100)

    # 连接电网络
    # 核心-聚合连接
    # 每个核心交换机连接所有Pod中特定聚合交换机
    for i, core in enumerate(core_e_nodes):
        # 计算核心交换机连接所有Pod中特定聚合交换机
        pod_offset=i //(k//2)
        aggr_in_pod=i % (k//2)+1
        for pod in range(1,pods+1):
            aggr_e_nodes = f"Aggr_E{pod}-{aggr_in_pod}"
            G.add_node(aggr_e_nodes, type='aggregation_electric')
            G.add_edge(core, aggr_e_nodes,
                       type='electrical',
                       capacity=10)  # 40Gbps核心-聚合链路
        # 连接所有Pod的第aggr_in-pod个聚合交换机
        # for aggr_id in range(i * 2 + 1, i * 2 + k+1):
        #     if aggr_id <= len(aggr_e_nodes):
        #         G.add_edge(core, f"{AGGR_ELECTRIC_PREFIX}{aggr_id}",
        #                    type='electrical', capacity=10)

    # 聚合-Edge交换机连接（每个聚合连接2个Edge）
    for pod in range(1, pods + 1):
        # 当前Pod的聚合交换机 (k//2个)
        aggr_e_nodes = [f"Aggr_E{pod}-{i}" for i in range(1, k // 2 + 1)]
        # 当前Pod的边缘交换机 (k//2个)
        edge_nodes = [f"Edge_E{pod}-{i}" for i in range(1, k // 2 + 1)]
        G.add_nodes_from(edge_nodes, type='edge_electric')

        # 每个聚合交换机连接本Pod所有边缘交换机
        for aggr in aggr_e_nodes:
            for edge in edge_nodes:
                G.add_edge(aggr, edge,
                           type='electrical',
                           capacity=10)  # 40Gbps聚合-边缘链路
    # for i, aggr in enumerate(aggr_e_nodes):
    #     for tor_id in range(i * 2 + 1, i * 2 + k//2+1):
    #         if tor_id <= k:
    #             G.add_edge(aggr, f"{EDGE_ELECTRIC_PREFIX}{tor_id}",
    #                        type='electrical',
    #                        capacity=10,
    #                        port_type='uplink')  # 标记为电上行链路

    # Edge-ToR连接（每个聚合连接1个ToR）
    for pod in range(1, pods + 1):
        # 当前的边缘交换机
        edge_nodes = [f"Edge_E{pod}-{i}" for i in range(1, k // 2 + 1)]
        # 当前Pod的ToR交换机
        tors_in_pod = [f"ToR{(pod - 1) * tor_per_pod + i}" for i in range(1, tor_per_pod + 1)]

        # 每个边缘交换机连接本Pod的k//2个ToR
        for edge_idx, edge in enumerate(edge_nodes):
            # 分配ToR的策略：轮询绑定
            connected_tors = tors_in_pod[edge_idx::k // 2]
            for tor in connected_tors:
                G.add_node(tor, type='tor')
                G.add_edge(edge, tor,
                           type='electrical',
                           capacity=10)  # 25Gbps边缘-ToR链路
    # for i, edge in enumerate(edge_nodes):
    #     for tor_id in range(i * 2 + 1, i * 2 + k//2):
    #         if tor_id <= k:
    #             G.add_edge(edge, f"{TOR_PREFIX}{tor_id}",
    #                        type='electrical',
    #                        capacity=10,
    #                        port_type='uplink')  # 标记为电上行链路

    # # 添加核心交换机之间的高速链路连接
    # for i in range(len(core_e_nodes)):
    #     for j in range(i + 1, len(core_e_nodes)):
    #         G.add_edge(core_e_nodes[i], core_e_nodes[j],
    #                    type='electrical', capacity=1600)  # 高速链路容量设为1000

    # ================== 光网络部分 ==================
    # 创建光交换机
    optical_switches = [f"{OPTICAL_SWITCH_PREFIX}{i}" for i in range(1, k*k*k//4 + 1)] # 创建光交换机的命名规则
    G.add_nodes_from(optical_switches, type='optical_switch', capacity=200) # 光交换机的容量设为1000

    # 连接光网络，让每个ToR连接所有光交换机
    for tor_id in range(1, k*k*k//4 + 1):
        tor_node = f"{TOR_PREFIX}{tor_id}"
        for sw_node in optical_switches:
            G.add_edge(tor_node, sw_node,
                       type='optical',
                       capacity=25,
                       port_type='optical')  # 标记为光端口

    return G


def visualize_topology(G):
    """可视化拓扑，分层排列节点"""
    # 定义节点颜色和形状
    node_colors = {
        'tor': 'lightblue',
        'edge_electric':'lightblue',
        'server': 'lightgreen',
        'core_electric': 'orange',
        'aggregation_electric': 'yellow',
        'optical_switch': 'pink',
        # 'aggregation_optical': 'purple'
    }
    node_shapes = {
        'tor': 's',  # 方形
        'edge_electric': 's',  # 方形
        'server': 'o',  # 圆形
        'core_electric': 'd',  # 菱形
        'aggregation_electric': '^',  # 三角形
        'optical_switch': 'p',  # 五边形
        # 'aggregation_optical': 'h'  # 六边形
    }

    # 定义边颜色
    edge_colors = {
        'electrical': 'gray',
        'optical': 'red'
    }

    # 手动指定节点位置
    pos = {}
    layer_spacing = 1.5  # 层间距
    node_spacing = 1.0  # 节点间距

    # 第一行：光交换机
    optical_switches = [n for n in G.nodes if G.nodes[n]['type'] == 'optical_switch']
    for i, node in enumerate(sorted(optical_switches)):
        pos[node] = (i * node_spacing, 5 * layer_spacing)


    # 第二行：ToR 交换机
    tor_switches = [n for n in G.nodes if G.nodes[n]['type'] == 'tor']
    for i, node in enumerate(sorted(tor_switches)):
        pos[node] = (i * node_spacing, 4 * layer_spacing)

    # 第三行：Edge 交换机
    edge_switches = [n for n in G.nodes if G.nodes[n]['type'] == 'edge_electric']
    for i, node in enumerate(sorted(edge_switches)):
        pos[node] = (i * node_spacing, 3 * layer_spacing)

    # 第三行：服务器
    servers = [n for n in G.nodes if G.nodes[n]['type'] == 'server']
    for i, node in enumerate(sorted(servers)):
        pos[node] = (i * node_spacing, 2 * layer_spacing)

    # 第四行：电聚合交换机
    aggr_electric = [n for n in G.nodes if G.nodes[n]['type'] == 'aggregation_electric']
    for i, node in enumerate(sorted(aggr_electric)):
        pos[node] = (i * node_spacing, 1 * layer_spacing)

    # 第五行：电核心交换机
    core_electric = [n for n in G.nodes if G.nodes[n]['type'] == 'core_electric']
    for i, node in enumerate(sorted(core_electric)):
        pos[node] = (i * node_spacing, 0 * layer_spacing)

    # 绘制节点
    plt.figure(figsize=(12, 8))
    for node_type, color in node_colors.items():
        nodes = [n for n in G.nodes if G.nodes[n]['type'] == node_type]
        nx.draw_networkx_nodes(G, pos, nodelist=nodes, node_color=color,
                               node_shape=node_shapes[node_type], label=node_type, node_size=300)

    # 绘制边
    for edge_type, color in edge_colors.items():
        edges = [(u, v) for u, v, d in G.edges(data=True) if d['type'] == edge_type]
        nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color=color, label=edge_type, width=2)

    # 绘制标签
    nx.draw_networkx_labels(G, pos, font_size=8, font_weight='bold')

    # 显示图例
    plt.legend(scatterpoints=1, frameon=False, title="Node Types", bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.title("Hybrid Topology with Shared ToR (Layered Layout)", fontsize=14)
    plt.tight_layout()
    plt.show()


# 示例用法
if __name__ == "__main__":
    network = build_hybrid_topology_shared_tor(k=4)
    # 检查缺失 type 的节点
    for node in network.nodes:
        print(node)
        if 'type' not in network.nodes[node]:
            print(f"错误：节点 {node} 缺少 'type' 属性！")
    # # 验证ToR连接
    tor1 = "ToR1"
    print(f"\n{tor1}的连接:")
    for neighbor in network.neighbors(tor1):
        edge_data = network.get_edge_data(tor1, neighbor)
        print(edge_data)
        # print(f" - {neighbor} ({edge_data['type']} link, port: {edge_data['port_type']})")


    # tor10 = "ToR10"
    # print(f"\n{tor10}的连接:")
    # for neighbor in network.neighbors(tor10):
    #     edge_data = network.get_edge_data(tor10, neighbor)
    #     print(f" - {neighbor} ({edge_data['type']} link, port: {edge_data['port_type']})")

    # 统计关键指标
    print("\n拓扑统计:")
    print(f"总节点数: {len(network.nodes)}")
    print(f"总连接数: {len(network.edges)}")
    print(f"电核心交换机: {len([n for n in network.nodes if 'core_electric' in network.nodes[n]['type']])}")
    print(f"电聚合交换机: {len([n for n in network.nodes if 'aggregation_electric' in network.nodes[n]['type']])}")
    print(f"电边缘交换机: {len([n for n in network.nodes if 'edge_electric' in network.nodes[n]['type']])}")
    print(f"光交换机: {len([n for n in network.nodes if 'optical_switch' in network.nodes[n]['type']])}")
    # print(f"光链路容量: {network.edges['core_electric1', 'optical_switch1']['capacity']}")
    # print(f"电链路容量: {sum([network.edges[u, v]['capacity'] for u, v in network.edges if 'electrical' in network.edges[u, v]['type']])}")
    # 输出每条链路的容量
    for u, v in network.edges():
        print(f"链路 {u}→{v} | 容量: {network.edges[u, v]['capacity']}")
    # 可视化拓扑
    visualize_topology(network)